#!/usr/bin/env python
#
# Copyright (C) 2011 Clement Lorteau <clem@lorteau.fr>	
#		
# This file is part of QuickScreenShots.		
# 	
# This program is free software; you can redistribute it and/or modify	
# it under the terms of the GNU General Public License as published by		
# the Free Software Foundation; version 3.	
# 	
# This program is distributed in the hope that it will be useful,		
# but WITHOUT ANY WARRANTY; without even the implied warranty of		
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the		
# GNU General Public License for more details.		
# 		
# You should have received a copy of the GNU General Public License		
# along with this program; if not, write to the Free Software		
# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301, USA.

import platform
from PIL import Image
from tempfile import NamedTemporaryFile
import os

def snap(showcursor = True):
	if platform.system() == 'Windows':
		return snap_win(showcursor)
	elif platform.system() == 'Darwin':
		return snap_osx(showcursor)
		
def snap_osx(showcursor = True):
	return (None, None)
    
def snap_win(showcursor = True):
    """Get the desktop DC and blit only the part where the active window is
       That's an easy way to capture the transparency effects of Vista/7/8
       (blitting the window's DC fails to get any transparency)"""
    #when using aero glass, getwindowrect doesn't return the expected
    #http://social.msdn.microsoft.com/Forums/en-US/windowsuidevelopment/thread/6c1c67a9-5548-4e9b-989f-c7dbac0b1375/
    #therefore our capture is a bit larger than the window
    #solution would be to use ctypes to call DwmGetWindowAttribute as mentioned here:
    #http://stackoverflow.com/questions/3192232/getwindowrect-too-small-on-windows-7
    import win32gui,  win32ui,  win32con, win32api, win32clipboard
    hwnd = win32gui.GetDesktopWindow()
    hwndfg = win32gui.GetForegroundWindow()
    l, t, r, b = win32gui.GetWindowRect(hwndfg)
    w = r - l
    h = b - t
     
    hwndDC = win32gui.GetWindowDC(hwnd)
    mfcDC  = win32ui.CreateDCFromHandle(hwndDC)
    saveDC = mfcDC.CreateCompatibleDC()

    saveBitMap = win32ui.CreateBitmap()
    saveBitMap.CreateCompatibleBitmap(mfcDC, w, h)
    saveDC.SelectObject(saveBitMap)
    saveDC.BitBlt((0, 0), (w, h),  mfcDC,  (l, t),  win32con.SRCCOPY)

    #add cursor
    if showcursor:
        curFlags, curH, (curX, curY) = win32gui.GetCursorInfo()
        curWidth = win32api.GetSystemMetrics(win32con.SM_CXCURSOR)
        curHeight = win32api.GetSystemMetrics(win32con.SM_CYCURSOR)
        relCurX = curX - l - curWidth/4
        relCurY = curY - t - curHeight/4
        saveDC.DrawIcon((relCurX, relCurY), curH)
       
    #load into PIL image
    """Avoid dealing with multiple color depth by writing the bitmap to disk
       then loading it from there"""
    tempfile = NamedTemporaryFile(prefix='qs_', delete=False)
    tempfile.close()
    saveBitMap.SaveBitmapFile(saveDC, tempfile.name)
    im = Image.open(tempfile.name)
    im.load()
    os.remove(tempfile.name)

    win32gui.DeleteObject(saveBitMap.GetHandle())
    saveDC.DeleteDC()
    mfcDC.DeleteDC()
    win32gui.ReleaseDC(hwnd, hwndDC)

    return (im, None)
    
    
if __name__ == "__main__":
    import time
    time.sleep(3)
    (im, bitmap) = snap()
    im.save('screenshot.png', 'PNG')